// server.ts
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import type { RequestHandlerExtra } from '@modelcontextprotocol/sdk/shared/protocol.js';
import {
  isInitializeRequest,
  type ServerNotification,
  type ServerRequest,
} from '@modelcontextprotocol/sdk/types.js';
import cors from 'cors';
import { randomUUID } from 'crypto';
import 'dotenv/config';
import express from 'express';
import { createServer } from 'node:http';
import z from 'zod';
import { logServer, registerTools, warnServer } from './tool';
import { registerWebSocket } from './ws';

const app = express();
app.use(express.json());
app.use(
  cors({
    origin: '*',
    credentials: true,
    allowedHeaders: ['Content-Type', 'mcp-session-id', 'mcp-protocol-version'],
    exposedHeaders: ['mcp-session-id'],
  })
);
logServer(`✅ DOUBAO_API_KEY:`, process.env.DOUBAO_API_KEY);
const transports: Record<string, StreamableHTTPServerTransport> = {};
const mcpServer = new McpServer({ name: 'example-server', version: '1.0.0' });
// 注册工具到 MCP Server
export const toolMap: Record<
  string,
  (input: any) => Promise<{
    content: (
      | { type: 'text'; text: string }
      | {
          type: 'resource_link';
          name: string;
          uri: string;
        }
    )[];
  }>
> = {};
export const registerToolToServer = <T extends z.ZodRawShape>(
  name: string,
  def: {
    title: string;
    description: string;
    inputSchema: T;
  },
  handler: (
    input: z.infer<z.ZodObject<T>>,
    session?: RequestHandlerExtra<ServerRequest, ServerNotification>
  ) => Promise<{
    content: (
      | { type: 'text'; text: string }
      | {
          type: 'resource_link';
          name: string;
          uri: string;
        }
    )[];
  }>
) => {
  try {
    mcpServer.registerTool(
      name,
      {
        title: def.title,
        description: def.description,
        inputSchema: def.inputSchema,
      },
      handler as any
    );
    console.log(`✅ 注册工具成功: ${name}`);
    toolMap[name] = handler as any;
  } catch (err) {
    console.error(`❌ 注册工具失败: ${name}`, err);
  }
};

// MCP 协议处理入口
app.post('/mcp', async (req, res) => {
  const sessionId = req.headers['mcp-session-id'] as string | undefined;
  logServer(`📨 MCP 请求: ${JSON.stringify(req.body, null, 2)}`);

  let transport: StreamableHTTPServerTransport;

  if (sessionId && transports[sessionId]) {
    logServer(`🔁 使用现有会话: ${sessionId}`);
    transport = transports[sessionId];
  } else if (!sessionId && isInitializeRequest(req.body)) {
    logServer(`🆕 初始化新会话`);
    transport = new StreamableHTTPServerTransport({
      sessionIdGenerator: () => randomUUID(),
      onsessioninitialized: (id) => {
        logServer(`✅ 会话初始化成功: ${id}`);
        transports[id] = transport;
      },
    });
    transport.onclose = () => {
      if (transport.sessionId) {
        logServer(`❎ 会话关闭: ${transport.sessionId}`);
        delete transports[transport.sessionId];
      }
    };
    await mcpServer.connect(transport);
  } else {
    warnServer(`⚠️ 无效会话 ID 或请求类型`);
    res.status(400).json({
      jsonrpc: '2.0',
      error: { code: -32000, message: 'No valid session ID provided' },
      id: null,
    });
    return;
  }

  await transport.handleRequest(req, res, req.body);
});

// 支持 GET / DELETE 保持 session
const handleSessionRequest = async (req: express.Request, res: express.Response) => {
  const sessionId = req.headers['mcp-session-id'] as string | undefined;
  if (!sessionId || !transports[sessionId]) {
    res.status(400).send('Invalid or missing session ID');
    return;
  }

  const transport = transports[sessionId];
  await transport.handleRequest(req, res);
};
registerTools();
app.get('/mcp', handleSessionRequest);
app.delete('/mcp', handleSessionRequest);

// 启动服务export async function startMcpServer(port = 4000) {
export async function startMcpServer(port = 4000) {
  const server = createServer(app);
  registerWebSocket(server); // ws 服务挂载到同一个 server
  return new Promise<void>((resolve) => {
    server.listen(port, () => {
      logServer(`✅ WS 服务已启动: ws://localhost:${port}`);
      logServer(`✅ MCP server listening on http://localhost:${port}/mcp`);
      resolve();
    });
  });
}
